from flask import Flask, render_template, request, jsonify, send_file
from langchain_community.vectorstores import Chroma
from langchain_openai import AzureChatOpenAI, AzureOpenAIEmbeddings
from langchain.prompts import PromptTemplate
from docx import Document
import os, datetime, io

# ==== Flask 初始化 ====
app = Flask(__name__)

# ==== Azure 配置（保持不变） ====
os.environ["AZURE_OPENAI_API_KEY"] = "EVlHLUCSdiDb0yVKwkm6USjrqshc2hZxyoLD5EI7PBzl3coU7ApTJQQJ99BCACYeBjFXJ3w3AAABACOGVmn4"
os.environ["AZURE_OPENAI_ENDPOINT"] = "https://eastus-0303.openai.azure.com/"
os.environ["AZURE_OPENAI_API_VERSION"] = "2024-08-01-preview"

embeddings = AzureOpenAIEmbeddings(
    deployment="text-embedding-3-large-1",
    model="text-embedding-3-large",
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version="2023-05-15"
)
db = Chroma(persist_directory="company_db", embedding_function=embeddings)

llm = AzureChatOpenAI(
    deployment_name="gpt-4o-1",
    model="gpt-4o",
    api_key=os.environ["AZURE_OPENAI_API_KEY"],
    azure_endpoint=os.environ["AZURE_OPENAI_ENDPOINT"],
    api_version=os.environ["AZURE_OPENAI_API_VERSION"],
    temperature=0
)

prompt = PromptTemplate(
    input_variables=["context", "question"],
    template="""你是一名经验丰富的售前技术专家。  
你的任务是基于以下资料，帮助销售人员快速理解客户的需求和痛点，并结合公司现有的产品、服务对象，指出可能的合作点和面。  

请遵循以下规则：  
1. 回答时要结构化，包含：
   - 客户需求/痛点  
   - 合作契机
   - 智用与合作公司的异同点
   - 公司能够提供的价值或解决方案  
   - 潜在合作点和面  
   - 下一步建议  
2. 如果资料中没有相关信息，请明确说明“未找到相关资料”，不要编造内容。  
3. 输出时要简洁清晰，便于销售快速把握重点。  

资料内容：  
{context}

用户问题：  
{question}

请基于以上信息给出专业回答："""
)

# ==== 对话历史 ====
conversation_history = []

# ==== 路由 ====
@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/ask', methods=['POST'])
def ask():
    question = request.json.get('question', '')
    docs = db.similarity_search(question, k=3)
    context = "\n\n".join([doc.page_content for doc in docs])
    final_prompt = prompt.format(context=context, question=question)
    response = llm.invoke(final_prompt)
    answer = response.content
    conversation_history.append({'question': question, 'answer': answer})
    return jsonify({'answer': answer})

@app.route('/api/export', methods=['GET'])
def export():
    if not conversation_history:
        return jsonify({'error': '无对话记录'}), 400
    doc = Document()
    doc.add_heading("RAG 问答记录", level=1)
    for item in conversation_history:
        doc.add_heading(f"问题：{item['question']}", level=2)
        doc.add_paragraph(item['answer'])
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return send_file(
        buffer,
        as_attachment=True,
        download_name=f"rag_qa_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}.docx"
    )

if __name__ == '__main__':
    app.run(debug=True)